import numpy as np
import nibabel as nib
import SimpleITK as sitk

from .utils import *


class DataLoader:
    """
    load image and labels and randomly generate patches.

    Args:
        file_path: file path of image data
        platform: deep learning platform used
        shape: size of the training data

    Attributes:
        testing_percentage: percentage of testing data
        LUT: lookup table for vertebraID corresponding to imageID
    """

    def __init__(self, file_path, platform='tensorflow', shape=[96, 96, 80]):
        self.__platform = platform
        self.testing_percentage = 0.2
        self.shrink_size = 3
        self.shape = shape
        self.LUT = np.zeros([2, 0])

        self.__prepare_data(file_path)

    def __get_info(self, label):
        min_label = label[label != 0].min()
        vrange = [min_label, label.max()]
        length = vrange[1] - vrange[0] + 1
        bounding_boxs = np.zeros([length, 6])

        for i in range(vrange[0], vrange[1] + 1):
            indexs = np.argwhere(label == i)
            bounding_box = [[indexs[:, axis].min(), indexs[:, axis].max()]
                            for axis in range(3)]
            bounding_boxs[i - vrange[0]] = np.hstack(bounding_box)

        index = 0 if self.LUT.shape[1] == 0 else self.LUT[1, -1] + 1
        print("Image %-3d : vertebrae %-2d - %-2d" %
              (index, vrange[0], vrange[1]))
        self.LUT = np.hstack(
            [self.LUT, [list(range(length)), np.ones(length) * index]])
        return {'range': vrange, 'bboxs': bounding_boxs}

    def __prepare_data(self, file_path):
        import os
        import random

        image_path = os.path.join(file_path, "image")
        label_path = os.path.join(file_path, "label")
        image_list = os.listdir(image_path)
        random.shuffle(image_list)

        self.label_info = []
        self.testing_data = []
        self.training_data = []

        self.num_testing = int(len(image_list) * self.testing_percentage)
        self.num_training = len(image_list) - self.num_testing

        for i, image_name in enumerate(image_list):
            image = nib.load(os.path.join(
                image_path, image_name)).get_data() + 1024  # input image should be in range -1024 ~ 3071
            image = np.pad(
                image, ((0, 0), (0, 0), (0, self.shape[2])), 'constant')

            label_name = image_name.replace('.nii.gz', '_label.nii.gz')
            label = nib.load(os.path.join(label_path, label_name)).get_data()
            label = np.pad(
                label, ((0, 0), (0, 0), (0, self.shape[2])), 'constant').astype(np.uint8)

            if i < self.num_training:
                info = self.__get_info(label)
                self.label_info.append(info)
                self.training_data.append([image, label])
            else:
                self.testing_data.append([image, label])
                print("Testing image :", image_name)

        self.LUT = self.LUT.astype(np.int16)

    # randomly generate samples for segmentation
    def get_data(self, iteration, batch_size=1):
        batch = np.zeros(np.hstack([4, batch_size, self.shape, 1]))
        for n in range(batch_size):
            batch[:, n, ..., 0] = self.random_select_patch(iteration)

        if self.__platform.lower() == 'pytorch':
            raise NotImplementedError
            # batch = np.transpose(batch, (0, 1, 5, 2, 3, 4))

        batch = [np.squeeze(b, 0) for b in np.split(batch, 4, 0)]
        return [np.concatenate([batch[0], batch[2]], -1), batch[1], batch[3]]

    def get_testdata(self, batch_size=1):
        patch_shape = np.hstack([batch_size, self.shape, 1])
        patch_images = np.zeros(patch_shape)
        patch_labels = np.zeros(patch_shape)
        patch_memorys = np.zeros(patch_shape)

        for i in range(batch_size):
            imageID = rand(self.num_testing)
            image, label = self.testing_data[imageID]

            vrange = [label[label != 0].min(), label.max()]
            vid = rand(vrange[0], vrange[1])

            indexs = np.argwhere(label == vid)
            bounding_box = [[indexs[:, axis].min(), indexs[:, axis].max()]
                            for axis in range(3)]

            start = [bounding_box[i][1] - self.shape[i] + 1 for i in range(3)]
            end = [bounding_box[i][0] for i in range(3)]

            x = np.clip(rand(start[0], end[0]), 0, image.shape[0] - self.shape[0])
            y = np.clip(rand(start[1], end[1]), 0, image.shape[1] - self.shape[1])
            z = np.clip(rand(start[2], end[2]), 0, image.shape[2])
            patch_image = slide(image, [x, y, z], self.shape).copy()
            patch_label = slide(label, [x, y, z], self.shape).copy()

            patch_memory = (patch_label > 0) * 1 - (patch_label >= vid) * 1
            patch_label = (patch_label == vid) * 1

            patch_images[i, ..., 0] = normalize(patch_image)
            patch_labels[i, ..., 0] = patch_label
            patch_memorys[i, ..., 0] = patch_memory

        return [np.concatenate([patch_images, patch_memorys], -1), patch_labels]

    def random_select_empty_patch(self, augment_type):
        imageID = rand(self.num_training)
        image, label = self.training_data[imageID]

        x = rand(self.shape[0], image.shape[0]) - self.shape[0]
        y = rand(self.shape[1], image.shape[1]) - self.shape[1]
        z = rand(image.shape[2] - 2 * self.shape[2] + 1)
        # z = rand(0, min(60, image.shape[2] - self.shape[2]))
        patch_image = slide(image, [x, y, z], self.shape).copy()
        patch_label = slide(label, [x, y, z], self.shape).copy()

        l = patch_label[patch_label != 0].min() if patch_label.sum() > 0 else 0
        if augment_type % 2 and l > 0 and (patch_label == l).sum() > 1000:
            patch_memory = np.zeros(self.shape)
            patch_label = (patch_label == l) * 1
        else:
            patch_memory = (patch_label > 0) * 1
            patch_label[::] = 0

        return self.__data_augment(patch_image, patch_label, patch_memory, augment_type)

    def random_select_patch(self, augment_type):
        vid = rand(self.LUT.shape[1])
        imageID = self.LUT[1][vid]
        image, label = self.training_data[imageID]
        vertebraID = self.LUT[0][vid]

        info = self.label_info[imageID]['bboxs'][vertebraID]
        vertebraID += self.label_info[imageID]['range'][0]

        start = [info[i * 2 + 1] - self.shape[i] +
                 1 - self.shrink_size for i in range(3)]
        end = [info[i * 2] + self.shrink_size for i in range(3)]

        x = np.clip(rand(start[0], end[0]), 0, image.shape[0] - self.shape[0])
        y = np.clip(rand(start[1], end[1]), 0, image.shape[1] - self.shape[1])
        z = np.clip(rand(start[2], end[2]), 0, image.shape[2])

        patch_image = slide(image, [x, y, z], self.shape).copy()
        patch_label = slide(label, [x, y, z], self.shape).copy()
        patch_memory = (patch_label > 0) * \
            (patch_label < vertebraID).astype(np.uint8)
        patch_label = (patch_label == vertebraID).astype(np.uint8)
        return self.__data_augment(patch_image, patch_label, patch_memory, augment_type)

    def __data_augment(self, image, label, memory, ptype):
        if ptype % 5 == 1:
            image = gauss_smoothing(image)
        elif ptype % 5 == 2:
            image = gauss_noising(image)
        elif ptype % 5 == 3:
            image, label, memory = flip_image([image, label, memory], rand(2))
        elif ptype % 5 == 4:
            #image, label, memory = rotate_image([image, label, memory])
            pass

        distance = distance_transform(label)
        # label = label_smoothing(label)
        return [normalize(image), label, memory, distance]

'''
class VertebraeSequence(tf.keras.utils.Sequence):
    """Wrapper of DataLoader"""
    def __init__(self, file_path, platform='tensorflow', shape=[96, 96, 80]):
        self.loader = DataLoader(file_path, platform, shape)

    def __len__(self):
        return 10000

    def __getitem__(self, idx):
        image, label, memory, distance = self.loader.get_data(idx, 4)
        return np.concatenate([image, memory], -1), label
'''
